from math import log
import operator
'''
这个程序用来自己实现决策树
'''
def calcShannonEnt(dataset):
    numEntries=len(dataset)
    labelCounts={}
    for featvec in dataset:
        currentLabel=featvec[-1]
        if currentLabel not in labelCounts:
            labelCounts[currentLabel]=0
        labelCounts[currentLabel]+=1
    shannonEnt=0.0
    for key in labelCounts:
        prob=float(labelCounts[key])/numEntries
        shannonEnt-=prob*log(prob,2)
    return shannonEnt
# 根据给定的特征来划分数据集
def spilitDataset(dataset,axis,value):
    retDataset=[]
    for featVec in dataset:
        if featVec[axis] ==value:
            reducedFeatVec=featVec[:axis]
            reducedFeatVec.extend(featVec[axis+1:])
            retDataset.append(reducedFeatVec)
    return retDataset
def chooseBestFeatureToSplit(daset):
    numFeatures=len(dataset[0])-1
    baseEntropy=calcShannonEnt(dataset)
    bestInfogain=0.0
    bestFeature=-1
    for i in range(numFeatures):
        featlist=[example[i] for example in dataset]
        uniqueVals=set(featlist)
        newEntropy=0.0
        for value in uniqueVals:
            subDataSet=spilitDataset(dataset,i,value)
            prob=len(subDataSet)/float(len(dataset))
            newEntropy+=prob*calcShannonEnt(subDataSet)
        infogain=baseEntropy-newEntropy
        if(infogain>bestInfogain):
            bestInfogain=infogain
            bestFeature=i
    return bestFeature
def majorityCnt(classlist):
    classCount={}
    for vote in classlist:
        if vote not in classCount:
            classCount[vote]=0
        classCount+=1
    sortedClasscount=sorted(classCount.items(),key=operator.itemgetter(1),reverse=True)
    return sortedClasscount[0][0]


def creatTree(dataset,labels):
    classlist=[example[-1] for example in dataset]
    if classlist.count(classlist[0])==len(classlist):
            return classlist[0]
    if len(dataset[0])==1:
        return majorityCnt(classlist)
    bestFeat=chooseBestFeatureToSplit(dataset)
    bestFeatLabel=labels[bestFeat]
    myTree={bestFeatLabel:{}}
    del(labels[bestFeat])
    featValues=[example[bestFeat] for example in dataset]
    uniqueValues=set(featValues)
    for value in uniqueValues:
        subLabels=labels[:]
        myTree[bestFeatLabel][value]=creatTree(spilitDataset(dataset,bestFeat,value),subLabels)
    return myTree


if __name__ == '__main__':
    dataset=[[1,1,'yes'],
             [1,1,'yes'],
             [1,0,'no'],
             [0,1,'no'],
             [0,1,'no']
             ]
    labels=['no surfacing','flippers']
    # # 计算给定数据的香农熵
    # print(calcShannonEnt(dataset))
    # # 计算那个特征的信息增益最大
    # print(chooseBestFeatureToSplit(dataset))
    print(creatTree(dataset,labels))